# -*- coding: utf-8 -*-
import numpy as np
import cv2
import sys
import os, shutil
import time
import datetime
from imutils.object_detection import non_max_suppression
import imutils  # 安装库pip install imutils ；pip install --upgrade imutils更新版本大于v0.3.1
from imutils.video import FileVideoStream


class bug_rect_info:
    def __init__(self, x, y, w, h):
        self.x = x
        self.y = y
        self.w = w
        self.h = h
        self.count = 1


class people_detect:

    def __init__(self):
        # 已经训练好的分类器，检测佩戴安全头的人员
        self.face_cascade = cv2.CascadeClassifier(os.path.join(sys.path[0], 'cascade-v2.5.xml'))

        # 检测行人
        self.hog = cv2.HOGDescriptor()  # 初始化方向梯度直方图描述子
        # 设置支持向量机(Support Vector Machine)使得它成为一个预先训练好了的行人检测器
        self.hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
        self.font = cv2.FONT_HERSHEY_SIMPLEX

        self.IsExit = False

        self.dir_temp = os.path.join(sys.path[0], "temp")
        self.dir_people = os.path.join(sys.path[0], "people")
        self.dir_hat = os.path.join(sys.path[0], "hat")
        self.dir_empty = os.path.join(sys.path[0], "empty")
        self.dir_head = os.path.join(sys.path[0], "head")

        self.bug_list = []

        # 创建临时目录
        if not os.path.exists(self.dir_temp): os.makedirs(self.dir_temp)
        if not os.path.exists(self.dir_people): os.makedirs(self.dir_people)
        if not os.path.exists(self.dir_hat): os.makedirs(self.dir_hat)
        if not os.path.exists(self.dir_temp): os.makedirs(self.dir_temp)
        if not os.path.exists(self.dir_empty): os.makedirs(self.dir_empty)
        if not os.path.exists(self.dir_head): os.makedirs(self.dir_head)

    # 获取行人
    def get_people(self, frame, small_frame):
        # 检测行人
        (rects, weights) = self.hog.detectMultiScale(small_frame, winStride=(8, 8), padding=(8, 8), scale=1.15)

        rects = np.array([[x, y, x + w, y + h] for (x, y, w, h) in rects])
        index = 0

        offset_x = float(frame.shape[0]) / float(small_frame.shape[0])
        offset_y = float(frame.shape[1]) / float(small_frame.shape[1])

        print(offset_x, offset_y)

        head = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
        for (x, y, w, h) in rects:
            if weights[index] <= 0.68: continue

            # 检测安全帽
            face = self.get_safe_hat(small_frame[y:h, x:w])
            color = (0, 0, 255)
            if not face == None:
                color = (0, 255, 0)

                cv2.rectangle(frame, (int((x + face[0]) * offset_x), int((y + face[1]) * offset_y)),
                              (int((x + face[2]) * offset_x), int((y + face[3]) * offset_y)), (0, 255, 0), 2)

            cv2.rectangle(frame, (int(x * offset_x), int(y * offset_y)), (int(w * offset_x), int(h * offset_y)), color,
                          2)

            index = index + 1
            cv2.imwrite(os.path.join(self.dir_temp, head + str(index) + ".jpg"), frame)

    # 检测安全帽
    def get_safe_hat(self, people_frame):
        head_frame = people_frame[0: people_frame.shape[0] / 4, 0: people_frame.shape[1]];
        gray = cv2.cvtColor(head_frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.01, 2)

        face_rects = np.array([[f_x, f_y, f_x + f_w, f_y + f_h] for (f_x, f_y, f_w, f_h) in faces])
        face_pick = non_max_suppression(face_rects, probs=None, overlapThresh=0.65)

        face = (sys.maxint, sys.maxint, sys.maxint, sys.maxint)
        for (x, y, w, h) in face_pick:
            if y < face[1]:  # and f[1] < rect[1] / 3:
                face = (x, y, w, h)

        if face[0] == sys.maxint:
            face = None;

        return face

    # 处理本地临时存储的图片，检测行人
    def local_image_file_people(self):
        list = os.listdir(self.dir_temp)
        capture_index = 0
        time.sleep(1)  # 拿到文件列表后等一秒，避免文件写入冲突
        for i in range(0, len(list)):
            path = os.path.join(self.dir_temp, list[i])
            print(path)
            frame = cv2.imread(path)
            (rects, weights) = self.hog.detectMultiScale(frame, winStride=(8, 8), padding=(8, 8), scale=1.15)
            rects = np.array([[x, y, x + w, y + h] for (x, y, w, h) in rects])
            pick = non_max_suppression(rects, probs=None, overlapThresh=0.65)

            # 没找到行人
            if len(pick) == 0:
                if os.path.exists(path):
                    head = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
                    shutil.move(path, os.path.join(self.dir_empty, head + str(i) + ".jpg"))
                    continue;

            # 保存行人信息
            for (x, y, w, h) in pick:
                flag = False

                # 移动范围小于5的，不记录
                for item in self.bug_list:
                    if abs(item.x - x) <= 2 and abs(item.y - y) <= 2:
                        flag = True
                        item.count = item.count + 1
                        print("step <= 2")
                        print(x, y, w, h)
                        break
                if flag:
                    if os.path.exists(path):
                        os.remove(path)
                    continue
                head = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
                people = frame[y: h, x: w];

                newpath = os.path.join(self.dir_people, head + str(capture_index) + ".jpg")
                cv2.imwrite(newpath, people)
                capture_index = capture_index + 1
                if os.path.exists(path):
                    os.remove(path)

                self.bug_list.append(bug_rect_info(x, y, w, h))
            time.sleep(0.1)

    # 处理本地临时存储的图片，检测安全帽
    def local_image_file_self_hat(self):
        list = os.listdir(self.dir_people)
        capture_index = 0
        time.sleep(1)
        for i in range(0, len(list)):
            path = os.path.join(self.dir_people, list[i])
            frame = cv2.imread(path)
            title = frame[0: frame.shape[0] / 4, 0: frame.shape[1]];
            head = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')

            gray = cv2.cvtColor(title, cv2.COLOR_BGR2GRAY)
            faces = self.face_cascade.detectMultiScale(gray, 1.01, 2)

            if len(faces) == 0:
                if os.path.exists(path):
                    shutil.move(path, os.path.join(self.dir_head, head + str(i) + ".jpg"))
                    continue;

            face_rects = np.array([[f_x, f_y, f_x + f_w, f_y + f_h] for (f_x, f_y, f_w, f_h) in faces])
            face_pick = non_max_suppression(face_rects, probs=None, overlapThresh=0.65)

            for (f_x, f_y, f_w, f_h) in face_pick:

                # cv2.rectangle(frame,(x,y),(x+w,y+h),(255,0,0),2)
                newpath = os.path.join(self.dir_hat, head + str(capture_index) + ".jpg")
                cv2.imwrite(newpath, title[f_y:f_h, f_x:f_w]);
                print
                newpath
                if os.path.exists(path):
                    os.remove(path)

                capture_index = capture_index + 1
            time.sleep(0.1)

    # 本地图片处理进程
    def people_process(self):
        while not self.IsExit:
            try:
                self.local_image_file_people();
                self.local_image_file_self_hat();

                # 找出前5个重复最多的位置，标记为异常位置。
                new_bug_list = []
                for i in range(5):
                    new_bug_list.append(bug_rect_info(0, 0, 0, 0))

                for item in self.bug_list:
                    min_index = 0
                    min_value = sys.maxint;
                    for i in range(len(new_bug_list)):
                        if new_bug_list[i].count < min_value:
                            min_value = new_bug_list[i].count;
                            min_index = i;
                    # print min_index, min_value;
                    new_bug_list[min_index] = item

                self.bug_list = new_bug_list


            except BaseException as err:
                print
                err
            finally:
                time.sleep(5)

    def main(self):
        # 工地摄像头
        # cap = cv2.VideoCapture(get_video_address.get_video_url())
        # 本地摄像头
        # cap = cv2.VideoCapture(0)
        # 本地文件
        #cap = FileVideoStream(os.path.join(sys.path[0], 'D:\img\safety\data\bc744__P720.mp4')).start()
        cap = cv2.VideoCapture('D:\\img\safety\data\P720.mp4')
        print(cap.isOpened())

        index = 0
        while True:
            try:
                ##head = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
                frame = cap.read()
                if frame is None:
                    break
                index = index + 1

                dt1 = datetime.datetime.now()
                # 修改图像的大小。 这个函数会根据图片的比例进行重新绘制大小
                # frame = imutils.resize(frame, width=min(1024, frame[1]))
                small_frame = imutils.resize(frame, width=min(480, frame.shape[1]))

                self.get_people(frame, small_frame)
                dt2 = datetime.datetime.now()
                small_frame = imutils.resize(small_frame, width=1024)
                print(dt2 - dt1)
                cv2.imshow('frame', frame)

                if cv2.waitKey(5) == 27:
                    break
            except AttributeError as err:
                print(err)
                #cap = self.getframes();

        self.IsExit = True


if __name__ == '__main__':
    test = people_detect()
    test.main()
